{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute this cell to install dependencies\n",
    "%pip install sf-hamilton[visualization]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example notebook [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dagworks-inc/hamilton/blob/main/examples/spark/world_of_warcraft/world_of_warcraft__spark_v1.ipynb) [![GitHub badge](https://img.shields.io/badge/github-view_source-2b3137?logo=github)](https://github.com/apache/hamilton/blob/main/examples/spark/world_of_warcraft/world_of_warcraft__spark_v1.ipynb)\n",
    "\n",
    "\n",
    "This notebooks shows how you can create multiple features that build/depend upon each other, using spark."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-05-22T22:27:16.952106Z",
     "start_time": "2023-05-22T22:27:15.116241Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:'PYARROW_IGNORE_TIMEZONE' environment variable was not set. It is required to set this environment variable to '1' in both driver and executor sides if you use pyarrow>=2.0.0. pandas-on-Spark will set it for you but it does not work if there is a Spark context already launched.\n"
     ]
    }
   ],
   "source": [
    "# Cell 1 - import the things you need\n",
    "import logging\n",
    "import sys\n",
    "import importlib\n",
    "\n",
    "import pyspark.sql as ps\n",
    "import pandas as pd\n",
    "\n",
    "from hamilton import driver\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-05-22T22:27:16.982107Z",
     "start_time": "2023-05-22T22:27:16.954791Z"
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Cell 2 - import modules to create part of the DAG from\n",
    "# We use the autoreload extension that comes with ipython to automatically reload modules when\n",
    "# the code in them changes.\n",
    "\n",
    "# import the jupyter extension\n",
    "%load_ext autoreload\n",
    "# set it to only reload the modules imported\n",
    "%autoreload 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:hamilton.telemetry:Note: Hamilton collects completely anonymous data about usage. This will help us improve Hamilton over time. See https://github.com/apache/hamilton#usage-analytics--data-privacy for details.\n"
     ]
    },
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 7.0.6 (20230106.0513)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"346pt\" height=\"332pt\"\n",
       " viewBox=\"0.00 0.00 345.53 332.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 328)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-328 341.53,-328 341.53,4 -4,4\"/>\n",
       "<!-- zone_flags -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>zone_flags</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"260.84\" cy=\"-162\" rx=\"48.99\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"260.84\" y=\"-158.3\" font-family=\"Times,serif\" font-size=\"14.00\">zone_flags</text>\n",
       "</g>\n",
       "<!-- zone_counts -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>zone_counts</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"178.84\" cy=\"-90\" rx=\"54.69\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"178.84\" y=\"-86.3\" font-family=\"Times,serif\" font-size=\"14.00\">zone_counts</text>\n",
       "</g>\n",
       "<!-- zone_flags&#45;&gt;zone_counts -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>zone_flags&#45;&gt;zone_counts</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M242.23,-145.12C231.69,-136.12 218.3,-124.68 206.54,-114.65\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"208.97,-112.12 199.1,-108.29 204.43,-117.44 208.97,-112.12\"/>\n",
       "</g>\n",
       "<!-- spark_session -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>spark_session</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"260.84\" cy=\"-306\" rx=\"59.59\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"260.84\" y=\"-302.3\" font-family=\"Times,serif\" font-size=\"14.00\">spark_session</text>\n",
       "</g>\n",
       "<!-- world_of_warcraft -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>world_of_warcraft</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"260.84\" cy=\"-234\" rx=\"76.89\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"260.84\" y=\"-230.3\" font-family=\"Times,serif\" font-size=\"14.00\">world_of_warcraft</text>\n",
       "</g>\n",
       "<!-- spark_session&#45;&gt;world_of_warcraft -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>spark_session&#45;&gt;world_of_warcraft</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M260.84,-287.7C260.84,-280.41 260.84,-271.73 260.84,-263.54\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"264.34,-263.62 260.84,-253.62 257.34,-263.62 264.34,-263.62\"/>\n",
       "</g>\n",
       "<!-- zone_likelihoods -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>zone_likelihoods</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"178.84\" cy=\"-18\" rx=\"70.69\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"178.84\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\">zone_likelihoods</text>\n",
       "</g>\n",
       "<!-- aggregation_level -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>aggregation_level</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" cx=\"96.84\" cy=\"-162\" rx=\"96.68\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"96.84\" y=\"-158.3\" font-family=\"Times,serif\" font-size=\"14.00\">Input: aggregation_level</text>\n",
       "</g>\n",
       "<!-- aggregation_level&#45;&gt;zone_counts -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>aggregation_level&#45;&gt;zone_counts</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M116.69,-144.05C127.01,-135.24 139.8,-124.33 151.09,-114.69\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"153.25,-117.44 158.59,-108.29 148.71,-112.12 153.25,-117.44\"/>\n",
       "</g>\n",
       "<!-- world_of_warcraft&#45;&gt;zone_flags -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>world_of_warcraft&#45;&gt;zone_flags</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M260.84,-215.7C260.84,-208.41 260.84,-199.73 260.84,-191.54\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"264.34,-191.62 260.84,-181.62 257.34,-191.62 264.34,-191.62\"/>\n",
       "</g>\n",
       "<!-- zone_counts&#45;&gt;zone_likelihoods -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>zone_counts&#45;&gt;zone_likelihoods</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M178.84,-71.7C178.84,-64.41 178.84,-55.73 178.84,-47.54\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"182.34,-47.62 178.84,-37.62 175.34,-47.62 182.34,-47.62\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x10535fbe0>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Cell 3 - Define the inputs, modules containing feature logic and the driver\n",
    "aggregation_level = \"avatarId\"\n",
    "zone_features = importlib.import_module(\"zone_features__spark_v1\")\n",
    "\n",
    "dr = driver.Driver({}, zone_features)  # can pass in multiple modules\n",
    "dr.display_all_functions(None)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting default log level to \"WARN\".\n",
      "To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23/09/05 10:57:33 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable\n",
      "23/09/05 10:57:33 WARN Utils: Service 'SparkUI' could not bind on port 4040. Attempting port 4041.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:hamilton.base:It appears no Pandas index type was detected (ignore this warning if you're using DASK for now.) Please check whether the dataframe created matches what what you expect to happen.\n",
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>avatarId</th>\n",
       "      <th>total_count</th>\n",
       "      <th>darkshore_count</th>\n",
       "      <th>durotar_count</th>\n",
       "      <th>durotar_likelihood</th>\n",
       "      <th>darkshore_likelihood</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>29</td>\n",
       "      <td>5597</td>\n",
       "      <td>3</td>\n",
       "      <td>101</td>\n",
       "      <td>0.018045</td>\n",
       "      <td>0.000536</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1806</td>\n",
       "      <td>477</td>\n",
       "      <td>0</td>\n",
       "      <td>4</td>\n",
       "      <td>0.008386</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>9233</td>\n",
       "      <td>1611</td>\n",
       "      <td>0</td>\n",
       "      <td>10</td>\n",
       "      <td>0.006207</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2040</td>\n",
       "      <td>748</td>\n",
       "      <td>0</td>\n",
       "      <td>16</td>\n",
       "      <td>0.021390</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>26</td>\n",
       "      <td>2103</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>0.019020</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>91059</th>\n",
       "      <td>90894</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>91060</th>\n",
       "      <td>90865</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>91061</th>\n",
       "      <td>90700</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>91062</th>\n",
       "      <td>90655</td>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>91063</th>\n",
       "      <td>90767</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>91064 rows × 6 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "       avatarId  total_count  darkshore_count  durotar_count  \\\n",
       "0            29         5597                3            101   \n",
       "1          1806          477                0              4   \n",
       "2          9233         1611                0             10   \n",
       "3          2040          748                0             16   \n",
       "4            26         2103                0             40   \n",
       "...         ...          ...              ...            ...   \n",
       "91059     90894            1                0              0   \n",
       "91060     90865            1                0              1   \n",
       "91061     90700            3                0              0   \n",
       "91062     90655            2                0              0   \n",
       "91063     90767            3                0              0   \n",
       "\n",
       "       durotar_likelihood  darkshore_likelihood  \n",
       "0                0.018045              0.000536  \n",
       "1                0.008386              0.000000  \n",
       "2                0.006207              0.000000  \n",
       "3                0.021390              0.000000  \n",
       "4                0.019020              0.000000  \n",
       "...                   ...                   ...  \n",
       "91059            0.000000              0.000000  \n",
       "91060            1.000000              0.000000  \n",
       "91061            0.000000              0.000000  \n",
       "91062            0.000000              0.000000  \n",
       "91063            0.000000              0.000000  \n",
       "\n",
       "[91064 rows x 6 columns]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Cell 4 - Execute the pipeline and display the result\n",
    "\n",
    "result = dr.execute([\"zone_likelihoods\"], inputs={\"aggregation_level\": aggregation_level})\n",
    "\n",
    "result.iloc[0,0].toPandas()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
